// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/reader.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/vector/vector_test_base.h>

#include <gtest/gtest.h>

namespace kumo::pollux::dwio::common {
namespace {

using namespace kumo::pollux::common;

class ReaderTest : public testing::Test, public pollux::VectorBuilder {
 protected:
  static void SetUpTestCase() {
    memory::MemoryManager::testingSetInstance({});
  }
};

TEST_F(ReaderTest, getOrCreateChild) {
  auto input = make_row_vector(
      {"c.0", "c.1"},
      {
          make_flat_vector<int64_t>({1, 2, 3, 4, 5}),
          make_flat_vector<int64_t>({2, 4, 6, 7, 8}),
      });

  common::ScanSpec spec("<root>");
  spec.addField("c.0", 0);
  // Create child from name.
  spec.getOrCreateChild("c.1")->setFilter(
      common::createBigintValues({2, 4, 6}, false));

  auto actual = RowReader::projectColumns(input, spec, nullptr);
  auto expected = make_row_vector({
      make_flat_vector<int64_t>({1, 2, 3}),
  });
  test::assertEqualVectors(expected, actual);

  // Create child from subfield.
  spec.getOrCreateChild(common::Subfield("c.1"))
      ->setFilter(common::createBigintValues({2, 4, 6}, false));
  POLLUX_ASSERT_USER_THROW(
      RowReader::projectColumns(input, spec, nullptr),
      "Field not found: c. Available fields are: c.0, c.1.");
}

TEST_F(ReaderTest, projectColumnsFilterStruct) {
  constexpr int kSize = 10;
  auto input = make_row_vector({
      make_flat_vector<int64_t>(kSize, melon::identity),
      make_row_vector({
          make_flat_vector<int64_t>(kSize, melon::identity),
      }),
  });
  common::ScanSpec spec("<root>");
  spec.addField("c0", 0);
  spec.getOrCreateChild(common::Subfield("c1.c0"))
      ->setFilter(common::createBigintValues({2, 4, 6}, false));
  auto actual = RowReader::projectColumns(input, spec, nullptr);
  auto expected = make_row_vector({
      make_flat_vector<int64_t>({2, 4, 6}),
  });
  test::assertEqualVectors(expected, actual);
}

TEST_F(ReaderTest, projectColumnsFilterArray) {
  constexpr int kSize = 10;
  auto input = make_row_vector({
      make_flat_vector<int64_t>(kSize, melon::identity),
      make_array_vector<int64_t>(
          kSize,
          [](auto) { return 1; },
          [](auto i) { return i; },
          [](auto i) { return i % 2 != 0; }),
  });
  common::ScanSpec spec("<root>");
  spec.addField("c0", 0);
  auto* c1 = spec.getOrCreateChild(common::Subfield("c1"));
  {
    SCOPED_TRACE("IS NULL");
    c1->setFilter(std::make_unique<common::IsNull>());
    auto actual = RowReader::projectColumns(input, spec, nullptr);
    auto expected = make_row_vector({
        make_flat_vector<int64_t>({1, 3, 5, 7, 9}),
    });
    test::assertEqualVectors(expected, actual);
  }
  {
    SCOPED_TRACE("IS NOT NULL");
    c1->setFilter(std::make_unique<common::IsNotNull>());
    auto actual = RowReader::projectColumns(input, spec, nullptr);
    auto expected = make_row_vector({
        make_flat_vector<int64_t>({0, 2, 4, 6, 8}),
    });
    test::assertEqualVectors(expected, actual);
  }
}

TEST_F(ReaderTest, projectColumnsMutation) {
  constexpr int kSize = 10;
  auto input = make_row_vector({make_flat_vector<int64_t>(kSize, melon::identity)});
  common::ScanSpec spec("<root>");
  spec.addAllChildFields(*input->type());
  std::vector<uint64_t> deleted(bits::nwords(kSize));
  bits::setBit(deleted.data(), 2);
  Mutation mutation;
  mutation.deletedRows = deleted.data();
  auto actual = RowReader::projectColumns(input, spec, &mutation);
  auto expected = make_row_vector({
      make_flat_vector<int64_t>({0, 1, 3, 4, 5, 6, 7, 8, 9}),
  });
  test::assertEqualVectors(expected, actual);
  random::setSeed(42);
  random::RandomSkipTracker randomSkip(0.5);
  mutation.randomSkip = &randomSkip;
  actual = RowReader::projectColumns(input, spec, &mutation);
#if MELON_HAVE_EXTRANDOM_SFMT19937
  expected = make_row_vector({
      make_flat_vector<int64_t>({0, 1, 3, 5, 6, 8}),
  });
#elif __APPLE__
  expected = make_row_vector({
      make_flat_vector<int64_t>({1, 5, 6, 7, 8, 9}),
  });
#else
  expected = make_row_vector({
      make_flat_vector<int64_t>({3, 4, 7, 9}),
  });
#endif
  test::assertEqualVectors(expected, actual);
}

} // namespace
} // namespace kumo::pollux::dwio::common
